{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import h36motion3d as datasets\n",
    "from utils_chico.PoseDataset import PoseDataset\n",
    "from utils_chico.data_utils import normal_actions, abnorm_actions, normal_actions_, abnorm_actions_\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.optim as optim\n",
    "import torch.autograd\n",
    "import torch\n",
    "import numpy as np\n",
    "from utils.loss_funcs import *\n",
    "from utils.data_utils import define_actions\n",
    "from utils.h36_3d_viz import visualize\n",
    "\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import warnings\n",
    "from typing import Callable, Any, Optional, List\n",
    "\n",
    "import torch\n",
    "from torch import Tensor\n",
    "from torch import nn\n",
    "\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datas = 'h36m' # 'h36m' or 'chico'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#FLAGS FOR LOADING THE DATASET\n",
    "if datas == 'h36m':\n",
    "    path = r'C:\\Users\\108978\\Desktop\\seq2seq_GCN\\h3.6m\\dataset'\n",
    "else:\n",
    "    path = r'data\\CHICO\\dataset'\n",
    "input_n=10 # number of frames to train on(default=10)\n",
    "output_n=25 # number of frames to predict on\n",
    "input_dim=3 # dimensions of the input coordinates(default=3)\n",
    "skip_rate=1 # # skip rate of frames for H3.6M (default=1) \n",
    "\n",
    "if datas == 'h36m':\n",
    "    joints_to_consider=22  #joints for  H3.6M(default=22 for coordinates)\n",
    "else:\n",
    "    joints_to_consider=15\n",
    "\n",
    "#FLAGS FOR THE MODEL\n",
    "tcnn_layers=4 # number of layers for the Temporal Convolution of the Decoder (default=4)\n",
    "tcnn_kernel_size=[3,3] # kernel for the T-CNN layers (default=[3,3])\n",
    "input_dim=3 # dimensions of the input coordinates(default=3)\n",
    "st_gcnn_dropout=0.1 #(default=0.1)\n",
    "tcnn_dropout=0.0  ##(default=0.0)\n",
    "\n",
    "\n",
    "#FLAGS FOR THE TRAINING\n",
    "mode='test' #choose either train or test mode\n",
    "if datas == 'h36m':\n",
    "    n_epochs=41 \n",
    "    batch_size=256\n",
    "    batch_size_test=8\n",
    "    lr=1e-02\n",
    "    use_scheduler=True # use MultiStepLR scheduler\n",
    "    milestones=[25,30,37]  # the epochs after which the learning rate is adjusted by gamma ########### SOTA [25,30,37] \n",
    "else:\n",
    "    n_epochs=61\n",
    "    batch_size=128\n",
    "    batch_size_test=8\n",
    "    lr=1e-01 # learning rate\n",
    "    use_scheduler=True # use MultiStepLR scheduler\n",
    "    milestones=[5,25, 30, 40]  # the epochs after which the learning rate is adjusted by gamma ########### SOTA [25,30,37] \n",
    "gamma=0.1 #gamma correction to the learning rate, after reaching the milestone epochs\n",
    "clip_grad=None # select max norm to clip gradients\n",
    "model_path= r'C:\\Users\\108978\\Desktop\\seq2seq_GCN\\checkpoints' # path to the model checkpoint file\n",
    "model_name = datas+'_3d_'+str(output_n)+'frames' #the model name to save/load\n",
    "actions_to_consider_test='all' # actions to test on.\n",
    "\n",
    "\n",
    "#FLAGS FOR THE VISUALIZATION\n",
    "actions_to_consider_viz='all' # actions to visualize\n",
    "visualize_from='test'\n",
    "n_viz=2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if datas == 'h36m':\n",
    "    if output_n==25:\n",
    "        print('here')\n",
    "        maskA = np.load(r'checkpoints\\h36m\\masks\\maskA_25fps.npy') \n",
    "        maskT = np.load(r'checkpoints\\h36m\\masks\\maskT_25fps.npy') \n",
    "    else:\n",
    "        maskA = np.load(r'checkpoints\\h36m\\masks\\maskA_10fps.npy') \n",
    "        maskT = np.load(r'checkpoints\\h36m\\masks\\maskT_10fps.npy') \n",
    "else:\n",
    "    maskA = np.load(r'checkpoints\\chico\\masks\\maskA_25fps.npy')\n",
    "    maskT = np.load(r'checkpoints\\chico\\masks\\maskT_25fps.npy')\n",
    "\n",
    "maskA = torch.tensor(maskA)\n",
    "maskT = torch.tensor(maskT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.SeSGCN_student import Model\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using device: %s'%device)\n",
    "\n",
    "model= Model(input_dim,input_n,\n",
    "                        output_n,st_gcnn_dropout,joints_to_consider,tcnn_layers,tcnn_kernel_size,tcnn_dropout).float().to(device)\n",
    "\n",
    "optimizer=optim.Adam(model.parameters(),lr=lr,weight_decay=1e-05)\n",
    "\n",
    "if use_scheduler:\n",
    "    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=gamma)\n",
    "\n",
    "print('total number of parameters of the network is: '+str(sum(p.numel() for p in model.parameters() if p.requires_grad)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(dataset,vald_dataset, maskA, maskT):\n",
    "  train_loss = []\n",
    "  val_loss = []\n",
    "  val_loss_best = 1000\n",
    "  #dataset = datasets.Datasets(path,input_n,output_n,skip_rate, split=0)\n",
    "  print('>>> Training dataset length: {:d}'.format(dataset.__len__()))\n",
    "  data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True)\n",
    "\n",
    "  #vald_dataset = datasets.Datasets(path,input_n,output_n,skip_rate, split=1)\n",
    "  print('>>> Validation dataset length: {:d}'.format(vald_dataset.__len__()))\n",
    "  vald_loader = DataLoader(vald_dataset, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True)\n",
    "\n",
    "  dim_used = np.array([6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 21, 22, 23, 24, 25,\n",
    "                    26, 27, 28, 29, 30, 31, 32, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,\n",
    "                    46, 47, 51, 52, 53, 54, 55, 56, 57, 58, 59, 63, 64, 65, 66, 67, 68,\n",
    "                    75, 76, 77, 78, 79, 80, 81, 82, 83, 87, 88, 89, 90, 91, 92])\n",
    "\n",
    "\n",
    "\n",
    "  for epoch in range(n_epochs-1):\n",
    "      running_loss=0\n",
    "      n=0\n",
    "      model.train()\n",
    "      for cnt,batch in enumerate(data_loader): \n",
    "          batch=batch.float().to(device)\n",
    "          batch_dim=batch.shape[0]\n",
    "          n+=batch_dim\n",
    "\n",
    "          if datas == 'h36m':\n",
    "            sequences_train=batch[:, 0:input_n, dim_used].view(-1,input_n,len(dim_used)//3,3).permute(0,3,1,2)\n",
    "            sequences_gt=batch[:, input_n:input_n+output_n, dim_used].view(-1,output_n,len(dim_used)//3,3)\n",
    "          else:\n",
    "            sequences_train=batch[:, :input_n,].view(-1,input_n,joints_to_consider,3).permute(0,3,1,2)\n",
    "            sequences_gt=batch[:, input_n:input_n+output_n].view(-1,output_n,joints_to_consider,3)#.permute(0,3,1,2)\n",
    "\n",
    "\n",
    "          optimizer.zero_grad() \n",
    "\n",
    "          sequences_predict=model(sequences_train, maskA, maskT).permute(0,1,3,2)\n",
    "          \n",
    "          loss=mpjpe_error(sequences_predict,sequences_gt)\n",
    "\n",
    "\n",
    "          if cnt % 200 == 0:\n",
    "            print('[%d, %5d]  training loss: %.3f' %(epoch + 1, cnt + 1, loss.item())) \n",
    "\n",
    "          loss.backward()  \n",
    "          if clip_grad is not None:\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(),clip_grad)\n",
    "\n",
    "          optimizer.step()\n",
    "          running_loss += loss*batch_dim\n",
    "\n",
    "      train_loss.append(running_loss.detach().cpu()/n)  \n",
    "      model.eval()\n",
    "      with torch.no_grad():\n",
    "          running_loss=0 \n",
    "          n=0\n",
    "          for cnt,batch in enumerate(vald_loader):\n",
    "              batch=batch.float().to(device)\n",
    "              batch_dim=batch.shape[0]\n",
    "              n+=batch_dim\n",
    "              \n",
    "              if datas == 'h36m':\n",
    "                sequences_train=batch[:, 0:input_n, dim_used].view(-1,input_n,len(dim_used)//3,3).permute(0,3,1,2)\n",
    "                sequences_gt=batch[:, input_n:input_n+output_n, dim_used].view(-1,output_n,len(dim_used)//3,3)\n",
    "              else:\n",
    "                sequences_train=batch[:, :input_n,].view(-1,input_n,joints_to_consider,3).permute(0,3,1,2)\n",
    "                sequences_gt=batch[:, input_n:input_n+output_n].view(-1,output_n,joints_to_consider,3)#.permute(0,3,1,2)\n",
    "\n",
    "              sequences_predict=model(sequences_train, maskA, maskT).permute(0,1,3,2)\n",
    "\n",
    "              loss=mpjpe_error(sequences_predict,sequences_gt)\n",
    "              if cnt % 200 == 0:\n",
    "                        print('[%d, %5d]  validation loss: %.3f' %(epoch + 1, cnt + 1, loss.item())) \n",
    "              running_loss+=loss*batch_dim\n",
    "          val_loss.append(running_loss.detach().cpu()/n)\n",
    "          if running_loss/n < val_loss_best:\n",
    "            val_loss_best = running_loss/n\n",
    "            torch.save(model.state_dict(),model_name+'_Student_best')\n",
    "\n",
    "      if use_scheduler:\n",
    "        scheduler.step()\n",
    "\n",
    "\n",
    "      if (epoch+1)%10==0:\n",
    "        print('----saving model-----')\n",
    "        torch.save(model.state_dict(),model_name+'_Student')\n",
    "        plt.figure(1)\n",
    "        plt.plot(train_loss, 'r', label='Train loss')\n",
    "        plt.plot(val_loss, 'g', label='Val loss')\n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "        \n",
    "      if (epoch+1)==n_epochs:\n",
    "        print('----saving model-----')\n",
    "        torch.save(model.state_dict(),model_name+'_Student')\n",
    "\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Data\n",
    "if datas == 'h36m':\n",
    "    dataset = datasets.Datasets(path,input_n,output_n,skip_rate, split=0)\n",
    "    vald_dataset = datasets.Datasets(path,input_n,output_n,skip_rate, split=1)\n",
    "else:\n",
    "    dataset = PoseDataset(path,'train',input_n, output_n, actions=normal_actions, win_stride=1)\n",
    "    vald_dataset = PoseDataset(path,'validation', input_n,output_n, actions=normal_actions, win_stride=1)\n",
    "\n",
    "train(dataset, vald_dataset, maskA.to(device), maskT.to(device))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if datas == 'h36m':\n",
    "  def test(ckpt_path=None,maskA=None, maskT=None):\n",
    "\n",
    "    model.load_state_dict(torch.load(ckpt_path))\n",
    "    print('model loaded')\n",
    "    model.eval()\n",
    "    accum_loss=0  \n",
    "\n",
    "    summ = 0\n",
    "    cc = 0\n",
    "\n",
    "    accum_loss_mao=0\n",
    "    n_batches=0 # number of batches for all the sequences\n",
    "    actions=define_actions(actions_to_consider_test)\n",
    "    dim_used = np.array([6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 21, 22, 23, 24, 25,\n",
    "                      26, 27, 28, 29, 30, 31, 32, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,\n",
    "                      46, 47, 51, 52, 53, 54, 55, 56, 57, 58, 59, 63, 64, 65, 66, 67, 68,\n",
    "                      75, 76, 77, 78, 79, 80, 81, 82, 83, 87, 88, 89, 90, 91, 92])\n",
    "    # joints at same loc\n",
    "    joint_to_ignore = np.array([16, 20, 23, 24, 28, 31])\n",
    "    index_to_ignore = np.concatenate((joint_to_ignore * 3, joint_to_ignore * 3 + 1, joint_to_ignore * 3 + 2))\n",
    "    joint_equal = np.array([13, 19, 22, 13, 27, 30])\n",
    "    index_to_equal = np.concatenate((joint_equal * 3, joint_equal * 3 + 1, joint_equal * 3 + 2))\n",
    "    totalll=0\n",
    "    kkk=0\n",
    "    for action in actions:\n",
    "      loss_mao = 0\n",
    "      running_loss=0\n",
    "      n=0\n",
    "      dataset_test = datasets.Datasets(path,input_n,output_n,skip_rate, split=2,actions=[action])\n",
    "      #print('>>> test action for sequences: {:d}'.format(dataset_test.__len__()))\n",
    "\n",
    "      test_loader = DataLoader(dataset_test, batch_size=batch_size_test, shuffle=False, num_workers=0, pin_memory=True)\n",
    "      for cnt,batch in enumerate(test_loader):\n",
    "        with torch.no_grad():\n",
    "\n",
    "          batch=batch.to(device)\n",
    "          batch_dim=batch.shape[0]\n",
    "          n+=batch_dim\n",
    "          \n",
    "          \n",
    "          all_joints_seq=batch.clone()[:, input_n:input_n+output_n,:]\n",
    "\n",
    "          sequences_train=batch[:, 0:input_n, dim_used].view(-1,input_n,len(dim_used)//3,3).permute(0,3,1,2)\n",
    "          sequences_gt=batch[:, input_n:input_n+output_n, :]\n",
    "\n",
    "          \n",
    "          sss = time.time()\n",
    "          if maskA is None:\n",
    "            sequences_predict = model(sequences_train)\n",
    "          else:\n",
    "            sequences_predict = model(sequences_train, maskA, maskT)\n",
    "          totalll += time.time()-sss\n",
    "          kkk += 1\n",
    "          sequences_predict=sequences_predict.permute(0,1,3,2).contiguous().view(-1,output_n,len(dim_used))\n",
    "\n",
    "\n",
    "\n",
    "          all_joints_seq[:,:,dim_used] = sequences_predict\n",
    "\n",
    "\n",
    "          all_joints_seq[:,:,index_to_ignore] = all_joints_seq[:,:,index_to_equal]\n",
    "\n",
    "          loss=mpjpe_error(all_joints_seq.view(-1,output_n,32,3),sequences_gt.view(-1,output_n,32,3))\n",
    "\n",
    "\n",
    "          mpjpe_p3d_h36 = torch.sum(torch.mean(torch.norm(sequences_gt.view(-1,output_n,32,3) - all_joints_seq.view(-1,output_n,32,3), dim=3), dim=2), dim=0)\n",
    "\n",
    "          running_loss+=loss*batch_dim\n",
    "          accum_loss+=loss*batch_dim\n",
    "          loss_mao += mpjpe_p3d_h36[-1]\n",
    "          #accum_loss_mao.append(np.array(mpjpe_p3d_h36[-1]))\n",
    "\n",
    "      #print('loss at test subject for action : '+str(action)+ ' is: '+ str(running_loss/n))\n",
    "      #print(str(action),': ', str(np.round((running_loss/n).item(),1)))\n",
    "      print(str(action),': ', str(np.round((loss_mao/n).item(),1)))\n",
    "      summ += loss_mao/n\n",
    "      cc += 1\n",
    "      n_batches+=n\n",
    "    #print('Average: '+str(np.round((accum_loss/n_batches).item(),2)))\n",
    "    #print('Average MAO: '+str(np.round((summ/cc).item(),3)))\n",
    "    print('Prediction time: ', totalll/kkk)\n",
    "\n",
    "else:\n",
    "  def test(ckpt_path,maskA=None,maskT=None, fps=25):\n",
    "    all_mpjpe = []\n",
    "    model.load_state_dict(torch.load(ckpt_path))\n",
    "    print('model loaded')\n",
    "    model.eval()\n",
    "    accum_loss=0 \n",
    "    real_accum = 0 \n",
    "    n_batches=0 # number of batches for all the sequences\n",
    "    actions=define_actions(actions_to_consider_test)\n",
    "    dim_used = np.array([6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 21, 22, 23, 24, 25,\n",
    "                      26, 27, 28, 29, 30, 31, 32, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,\n",
    "                      46, 47, 51, 52, 53, 54, 55, 56, 57, 58, 59, 63, 64, 65, 66, 67, 68,\n",
    "                      75, 76, 77, 78, 79, 80, 81, 82, 83, 87, 88, 89, 90, 91, 92])\n",
    "    # joints at same loc\n",
    "    joint_to_ignore = np.array([16, 20, 23, 24, 28, 31])\n",
    "    index_to_ignore = np.concatenate((joint_to_ignore * 3, joint_to_ignore * 3 + 1, joint_to_ignore * 3 + 2))\n",
    "    joint_equal = np.array([13, 19, 22, 13, 27, 30])\n",
    "    index_to_equal = np.concatenate((joint_equal * 3, joint_equal * 3 + 1, joint_equal * 3 + 2))\n",
    "    totalll=0\n",
    "    kkk=0\n",
    "    for action in normal_actions:\n",
    "      mpjpe_joi, mpjpe_ali = np.zeros([fps]),np.zeros([fps])\n",
    "      running_loss=0\n",
    "      running_loss_real=0\n",
    "      n=0\n",
    "      dataset_test = PoseDataset(path,'test',input_n,output_n, actions=[action], win_stride=1)\n",
    "      #print('>>> test action for sequences: {:d}'.format(dataset_test.__len__()))\n",
    "\n",
    "      test_loader = DataLoader(dataset_test, batch_size=batch_size_test, shuffle=False, num_workers=0, pin_memory=True)\n",
    "      for cnt,batch in enumerate(test_loader):\n",
    "        with torch.no_grad():\n",
    "\n",
    "          batch=batch.to(device)\n",
    "          batch_dim=batch.shape[0]\n",
    "          n+=batch_dim\n",
    "          \n",
    "          \n",
    "          all_joints_seq=batch.clone()[:, input_n:input_n+output_n,:]\n",
    "\n",
    "          sequences_train=batch[:, 0:input_n, :, :].permute(0,3,1,2).float()\n",
    "          sequences_gt=batch[:, input_n:input_n+output_n, :,:].float()\n",
    "\n",
    "          \n",
    "          sss = time.time()\n",
    "          sequences_predict = model(sequences_train,maskA,maskT)\n",
    "          totalll += time.time()-sss\n",
    "          kkk += 1\n",
    "          sequences_predict=sequences_predict.permute(0,1,3,2).contiguous()\n",
    "\n",
    "   \n",
    "          tmp_joi = torch.sum(torch.mean(torch.norm(sequences_predict - sequences_gt, dim=3), dim=2), dim=0)\n",
    "          mpjpe_joi += tmp_joi.cpu().data.numpy()\n",
    "\n",
    "\n",
    "          loss=mpjpe_error(sequences_predict,sequences_gt)\n",
    "          \n",
    "          running_loss+=loss*batch_dim\n",
    "          accum_loss+=loss*batch_dim\n",
    "          \n",
    "      print(str(action),': ', str((running_loss/n).item()))\n",
    "      mpjpe_joi = mpjpe_joi/n\n",
    "      out_print_frame = [9,fps-1]\n",
    "      res_dic = {\"mpjpe_joi\": mpjpe_joi[out_print_frame]}\n",
    "      print('Error at each output frame:\\n Frame number:{}\\n Error:{}'.format(out_print_frame,res_dic))\n",
    "      all_mpjpe.append(mpjpe_joi[out_print_frame])\n",
    "      n_batches+=n\n",
    "    print('Average: '+str((accum_loss/n_batches).item()))\n",
    "    print(np.array(all_mpjpe).mean(0))\n",
    "    print('Prediction time: ', totalll/kkk)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if datas == 'h36m':\n",
    "    if output_n==10:\n",
    "        test(r'checkpoints\\h36m\\h36_3d_10frames_Student', maskA.to(device), maskT.to(device))\n",
    "    else:\n",
    "        test(r'checkpoints\\h36m\\h36_3d_25frames_Student', maskA.to(device), maskT.to(device))\n",
    "else:\n",
    "    if output_n==10:\n",
    "        test(r'checkpoints\\chico\\chico_3d_10frames_Student', maskA.to(device), maskT.to(device), fps=output_n)\n",
    "    else:\n",
    "        test(r'checkpoints\\chico\\chico_3d_25frames_Student', maskA.to(device), maskT.to(device), fps=output_n)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ad2bdc8ecc057115af97d19610ffacc2b4e99fae6737bb82f5d7fb13d2f2c186"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
